package com.itcast.spark.baseTree

import org.apache.spark.ml.{Pipeline, PipelineModel}
import org.apache.spark.ml.classification.{DecisionTreeClassificationModel, DecisionTreeClassifier}
import org.apache.spark.ml.evaluation.{BinaryClassificationEvaluator, MulticlassClassificationEvaluator}
import org.apache.spark.ml.feature.VectorAssembler
import org.apache.spark.mllib.tree.model.DecisionTreeModel
import org.apache.spark.{SparkConf, SparkContext}
import org.apache.spark.sql.{DataFrame, Dataset, Row, SparkSession}

/**
 * DESC:流程
 */
object _03LoveStortyModel {
  def main(args: Array[String]): Unit = {
    //1-准备环境
    val conf: SparkConf = new SparkConf().setAppName("_03LoveStortyModel").setMaster("local[*]")
    val spark: SparkSession = SparkSession.builder().config(conf).getOrCreate()
    val sc: SparkContext = spark.sparkContext
    sc.setLogLevel("WARN")
    //2-准备数据源
    val dataDF: DataFrame = spark.read.format("csv").option("header", true).option("inferschema", true).load("./datasets/mldata/love.txt")
    //3-数据的基本信息的查看
    //dataDF.printSchema()
    //dataDF.show(false)
    val split: Array[Dataset[Row]] = dataDF.randomSplit(Array(0.8, 0.2), seed = 1234L)
    val trainingSet: Dataset[Row] = split(0)
    val testSet: Dataset[Row] = split(1)
    //4-特征工程-is_date,age,is_handsome,income,is_gongwuyuan
    val assembler: VectorAssembler = new VectorAssembler().setInputCols(Array("age", "is_handsome", "income", "is_gongwuyuan")).setOutputCol("features")
    //* 5-准备算法
    val decisionTreeClassifier: DecisionTreeClassifier = new DecisionTreeClassifier()
      .setFeaturesCol("features")
      .setLabelCol("is_date")
      .setProbabilityCol("probability")
      .setPredictionCol("prediction") //预测为0-1的数值，后续可以对应到底是那个花
      .setMaxDepth(5)
      .setImpurity("gini")
      .setRawPredictionCol("rawPrediction")
    //* 6-模型的超参数的校验
    //* 7-模型训练
    val pipeline: Pipeline = new Pipeline().setStages(Array(assembler, decisionTreeClassifier))
    val pipelineModel: PipelineModel = pipeline.fit(trainingSet)
    //* 8-模型预测
    val y_pred_train: DataFrame = pipelineModel.transform(trainingSet)
    val y_pred_test: DataFrame = pipelineModel.transform(testSet)
    // *在模型校验的时候我们使用一般都是
    //new MulticlassClassificationEvaluator()  这个是多分类的教研和过程，也可以包括二分类
    //专门的二分类的场景
    val evaluator: BinaryClassificationEvaluator = new BinaryClassificationEvaluator()
      .setLabelCol("is_date")
      .setRawPredictionCol("rawPrediction")
      .setMetricName("areaUnderROC")//areaUnderROC==AUC
    val auc_train: Double = evaluator.evaluate(y_pred_train)
    val auc_test: Double = evaluator.evaluate(y_pred_test)
    println("auc train result is:",auc_train)
    println("auc test result is:",auc_test)
    //* 9-模型保存
    //* 10-已经保存的模型实现预测
    println(pipelineModel.stages(1).asInstanceOf[DecisionTreeClassificationModel].toDebugString)
  }
}
